# 这个模型只适合英文的微调哈
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from transformers import Trainer, TrainingArguments, AdamW
from torch.utils.data import DataLoader
import torch
import json
from typing import Dict
from torch.nn import CrossEntropyLoss
# openai gpt2 可以在huggingface上看看
# 加载预训练好的基本模型，同时把模型放置到gpu上面
tokenizer = GPT2Tokenizer.from_pretrained('gpt2',local_files_only=True)  # downloaded in default
model = GPT2LMHeadModel.from_pretrained('gpt2',local_files_only=True)
device="cuda:0"
model.to(device)
# 微调数据集的构建
with open("./data.jsonl", 'r', encoding='utf-8') as f:
     data=f.readlines()
#大多数llm都需要设置padding的token如果没有设置会报错，由于vocab.json：https://hf-mirror.com/openai-community/gpt2/tree/main 中是有eos_tokens的所以使用这个     
print(tokenizer.eos_token)
tokenizer.pad_token = tokenizer.eos_token
class Dataset(torch.utils.data.Dataset):
    def __init__(self, data:list[Dict]):
        self.data = data
   
    def __getitem__(self, idx):
        mid=json.loads(self.data[idx])
        text = mid["Q"] +' '+ mid["A"]#中间的这个字符是因为在openai提供的仓库中vocab.json本来就有的哈
        inputs = tokenizer.encode_plus(text, 
        max_length=128,  # 设置最大序列长度
        padding='max_length',  # 进行填充以确保序列达到最大长度
        truncation=True,  # 截断超出最大长度的序列
        return_attention_mask=True,  # 返回注意力掩码
        return_tensors='pt')
        attention_mask = inputs['attention_mask']
        #inputs.to(device)
        return {
            'input_ids': inputs['input_ids'],
            'attention_mask':attention_mask,   #0表示发生填充的地方,1表示没有填充的地方
            'labels': inputs['input_ids']
        }
    
        """
        theory explain: different to pre_train, SFT is not predicting next words, instead, it's to give the same sequence
        Q<sep>A,
        """
        # trainer类的数据格式应该如何构建
    def __len__(self):
        return len(self.data)

dataset=Dataset(data)
training_args = TrainingArguments(
    output_dir="./SFTgpt",
    overwrite_output_dir=True,
    num_train_epochs=500,              
    per_device_train_batch_size=5,  
    per_device_eval_batch_size=3,
    warmup_steps=1000,
    weight_decay=0.01,
    learning_rate=2e-3,
    logging_dir='./logs',
)

# 微调模型
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
)

# 训练模型
trainer.train()
# 7. 保存模型
model.save_pretrained('./SFTgpt')
tokenizer.save_pretrained('./SFTgpt')

# RuntimeError: cannot pin 'torch.cuda.LongTensor' only dense CPU tensors can be pinned